import argparse
import os

parser = argparse.ArgumentParser()
parser.add_argument('-c','--cuda', type=str, default='0', 
                help='CUDA device(s) to use, e.g. "0" or "0,1"')
args, _ = parser.parse_known_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda



from utils.MODEL_CKP import FLUX

import torch
from diffusers import FluxPipeline
from transformers import PipelineQuantizationConfig

# 配置 4-bit 量化
quantization_config = PipelineQuantizationConfig(
    load_in_4bit=True,               # 核心：启用 4-bit 加载
    bnb_4bit_quant_type="nf4",       # 使用一种特殊的 4-bit 数据类型 NF4，效果更好
    bnb_4bit_compute_dtype=torch.float16,  # 计算时使用 float16 加速
    # bnb_4bit_use_double_quant=True, # 可选项：嵌套量化，能再压缩一点模型，但可能更慢
)

# 从 Hugging Face Hub 加载量化的 Flux 模型
# 确保模型ID是正确的，例如 "black-forest-labs/FLUX.1-dev"
pipe = FluxPipeline.from_pretrained(
    FLUX,
    quantization_config=quantization_config, # 传入量化配置
    device_map="auto",               # 自动将模型各部分分布到可用的 GPU 和 CPU 上
    torch_dtype=torch.float16,      # 整体以 float16 模式加载，节省显存
)

# 现在你可以像平常一样使用 pipeline 了！
# 注意：生成的第一张图可能会比较慢，因为需要编译内核。
prompt = "A serene landscape with a river flowing through a lush green forest."
image = pipe(
    prompt,
    num_inference_steps=25,
    guidance_scale=7.0,
    height=1024,
    width=1024,
).images[0]

image.save("tmp_optim.png")